from typing import Generic, TypeVar, Union, Dict, List, Optional, Type

from mg_app_framework import get_logger

from bson import ObjectId
from pydantic import BaseModel, ValidationError
from pymongo import InsertOne, UpdateOne, UpdateMany, ReplaceOne, DeleteOne, DeleteMany
from motor.core import AgnosticCollection

from ..utils.singleton import Singleton

T = TypeVar('T', BaseModel, BaseModel)
OperationType = Union[InsertOne, UpdateOne, UpdateMany, ReplaceOne, DeleteOne, DeleteMany]


class BaseModelDao(Generic[T], Singleton):
    """为了兼容pydantic库的兼容写法, 新增加的表都使用这个类来加强数据校验, 防止mongo数据库字段类型随意定"""

    def get_db(self) -> AgnosticCollection:
        raise NotImplemented('必须实现get_db方法')

    async def insert_one(self, insert_data: T):
        db = self.get_db()
        insert_val = insert_data.dict(by_alias=True)
        result = await db.insert_one(insert_val)
        insert_data.id = result.inserted_id
        return result.inserted_id

    async def insert_many(self, data: List[T]):
        db = self.get_db()
        insert_dict_list = []
        for d in data:
            insert_dict_list.append(d.dict(by_alias=True))

        result = await db.insert_many(insert_dict_list)
        for id_, d in zip(result.inserted_ids, data):
            d.id = id_

        return result.inserted_ids

    async def insert_many_raw(self, data: List[Dict]):
        db = self.get_db()
        return await db.insert_many(data)

    async def insert_or_update_by_id(self, id_: Union[str, ObjectId], data: T):
        """根据id来更新或者插入一条新的数据"""
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        return await db.replace_one(query, data.dict(by_alias=True), upsert=True)

    async def insert_or_update_by_query(self, query: Dict, data: T):
        """根据id来更新或者插入一条新的数据"""
        db = self.get_db()
        return await db.replace_one(query, data.dict(by_alias=True), upsert=True)

    async def update_by_id(self, id_: Union[str, ObjectId], update_data: Dict):
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }

        update_ = {
            '$set': update_data
        }
        return await db.update_one(query, update_)

    async def update_many_by_query(self, query: Dict, update_data: Dict):
        db = self.get_db()
        update_ = {
            '$set': update_data
        }
        return await db.update_many(query, update_)

    async def get_by_id(self, id_: Union[str, ObjectId]) -> Optional[T]:
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        res = await db.find_one(query)
        if res:
            obj = self.dict_to_model(res)
            await self.after_find(obj)
            return obj
        return

    async def get_by_ids(self, ids: List[ObjectId]):
        db = self.get_db()
        query = {'_id': {'$in': ids}}
        return await self.get_by_query(query)

    async def get_one_by_query(self, query: Dict, **kwargs) -> Optional[T]:
        db = self.get_db()

        res = await db.find_one(query, **kwargs)
        if res:
            obj = self.dict_to_model(res)
            await self.after_find(obj)
            return obj
        return

    async def get_by_query(self, query: Dict, **kwargs) -> List[T]:
        db = self.get_db()
        res = []
        async for item in db.find(query, **kwargs):
            obj = self.dict_to_model(item)
            await self.after_find(obj)
            res.append(obj)
        return res

    async def delete_by_id(self, id_: Union[str, ObjectId]):
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        return await db.delete_one(query)

    async def delete_many_by_query(self, query: Dict):
        db = self.get_db()
        return await db.delete_many(query)

    async def delete_by_ids(self, ids: List[ObjectId]):
        db = self.get_db()
        query = {
            '_id': {'$in': ids}
        }
        return await db.delete_many(query)

    async def bulk_write(self, data: List[OperationType]):
        db = self.get_db()
        return await db.bulk_write(data)

    async def after_find(self, data: T):
        """找到对应的模型, 做对应的操作"""
        pass

    def _get_generate_class(self) -> Type[T]:
        """返回T的类型"""
        base_cls = self.__orig_bases__[0]
        cls = base_cls.__args__[0]
        return cls

    def dict_to_model(self, dic: Dict) -> T:
        cls = self._get_generate_class()
        try:
            obj = cls(**dic)
        except ValidationError as e:
            get_logger().warning(f'mongo result to model error: {e.errors()}')
            raise
        return obj


